# rag-service/app.py
from fastapi import FastAPI
from pymilvus import connections, Collection
import redis
import json

app = FastAPI(title="SHEIN RAG Service")

# 连接Milvus和Redis
connections.connect(host="milvus", port="19530")
redis_client = redis.Redis(host="redis", port=6379, decode_responses=True)

# 加载集合
collection = Collection("product_knowledge")


@app.post("/v1/retrieve")
async def retrieve_knowledge(query: dict):
    """混合检索接口"""
    query_text = query["text"]
    top_k = query.get("top_k", 5)

    # 混合检索逻辑
    search_params = {
        "data": [query_text],
        "anns_field": "dense_vector",
        "param": {"metric_type": "IP", "params": {"nprobe": 32}},
        "limit": top_k,
    }

    results = collection.search(**search_params)

    return {
        "results": [
            {
                "content": hit.entity.get("content"),
                "score": hit.score,
                "metadata": hit.entity.get("metadata")
            }
            for hit in results[0]
        ]
    }